__author__ = 'Janos G. Hajagos'

#    Classes for generating WaveFront 3D graphics files for visualizing data
#    in a more realistic space.

import math
import numpy as np

class ObjGraphic(object):
    def __init__(self, group_name = None, object_name = None, smoothing = False):
        """Basic class to represent WaveFront obj object.  Objects created by this class
        are meant to be combined together by a writer to create a single file."""
        self.vertices = None
        self.faces = None
        self.normal_vectors = None
        self.has_normal_vectors = False
        self.has_normal_indices = False
        self.texture_coordinates = None
        self.has_texture_coordinates = False
        self.has_texture_indices = False
        self.group_name = group_name
        self.object_name = object_name
        self.material_object = None
        self.smoothing = smoothing
        self.normal_indices = None
        self.texture_indices = None

    def associate_material(self,mtl_object):
        self.material_object = mtl_object

    def add_vertices(self, list_of_vertices):
        self.vertices = np.matrix(list_of_vertices)

    def add_normal_vectors(self, list_of_normal_vector):
        self.has_normal_vectors = True
        self.normal_vectors = np.matrix(list_of_normal_vector)

    def add_texture_coordinates(self, list_of_texture_coordinates):
        self.has_texture_coordinates = True
        self.texture_coordinates = np.matrix(list_of_texture_coordinates)

    def add_face(self, vertices_indices, normal_indices=None, texture_indices=None):
        self.vertices_indices = np.matrix(vertices_indices)

        if normal_indices:
            self.has_normal_indices = True
            self.normal_indices = np.matrix(normal_indices)

        if texture_indices:
            self.has_texture_indices = True
            self.texture_indices = np.matrix(texture_indices)

    def add_only_faces(self,faces):
        self.faces = np.matrix(faces)

    def add_normal_indices(self,normal_indices):
        self.has_normal_indices = True
        self.normal_indices = np.matrix(normal_indices)

    def _convert_regular_quad_to_mesh(self, face):
        """Converts a quadrilateral side to Triangular face:
        [4, 3,
         1, 2]
        [[1,3,4],[2,3,4]]
        """

        if len(face) != 4:
            raise RuntimeError, ""
        return [[face[0],face[1],face[2]],[face[2],face[3],face[0]]]

    def convert_regular_quads_to_mesh(self,face_list):
        """Converts a list of quad forms to triangulated mesh"""
        new_face_list = []
        for face in face_list:
            triangles = self._convert_regular_quad_to_mesh(face)
            new_face_list.append(triangles[0])
            new_face_list.append(triangles[1])
        return new_face_list

    def compute_normals_from_mesh(self):
        normal_vector_list = []
        for face in self.vertices_indices:
            i1 = face[0,0]
            i2 = face[0,1]
            i3 = face[0,2]

            v1 = self.vertices[i1 - 1]
            v2 = self.vertices[i2 - 1]
            v3 = self.vertices[i3 - 1]
            normal_vector_list.append(normal_vector_from_triangular_face(v3,v1,v2,return_list = True)) # this is based on triangulation pattern and right handed system with counter-clockwise rotation

        return normal_vector_list

def normal_vector_from_triangular_face(a,b,c,return_list = False):
    "Computes a normalized vector for a triangular face"
    cp = np.cross(b-a,c-a)
    ncp = cp / math.sqrt(cp[0,0]**2 + cp[0,1]**2 + cp[0,2]**2)
    if return_list:
        return [ncp[0,0],ncp[0,1],ncp[0,2]]
    else:
        return ncp

class classObjGeometryReader(object):
    """Ignores"""
    def __init__(self):
        pass
    def read_file(self):
        pass
    def read_string(self):
        pass

class MtlMaterial(object):
    """Class for representing MTL files the Wavefront compatible materials"""
    def __init__(self,name):
        self.name = name
        self.ambient_color = None # Ka
        self.diffuse_color = None #Kd
        self.specular_color = None # Ks
        self.specular_coefficient = None #Ns
        self.transparent = None # d or Tr range from 0 to 1
        self.illumination = None
        self.map_ambient = None # map_Ka           # the ambient texture map
        self.map_diffuse  = None# map_Kd filename.tga           # the diffuse texture map (most of the time, it will
        # be the same as the ambient texture map)
        self.map_specular = None # map_Ks lenna.tga           # specular color texture map
        self.map_specular_highlight = None #map_Ns filename.tga      # specular highlight component
        self.map_alpha_texture = None #map_d  filename.tga      # the alpha texture map
        self.map_bump = None #map_bump filename.tga

    def set_illumination(self, illumination):
        """
        0. Color on and Ambient off
        1. Color on and Ambient on
        2. Highlight on
        3. Reflection on and Ray trace on
        4. Transparency: Glass on, Reflection: Ray trace on
        5. Reflection: Fresnel on and Ray trace on
        6. Transparency: Refraction on, Reflection: Fresnel off and Ray trace on
        7. Transparency: Refraction on, Reflection: Fresnel on and Ray trace on
        8. Reflection on and Ray trace off
        9. Transparency: Glass on, Reflection: Ray trace off
        10. Casts shadows onto invisible surfaces

        Source: http://en.wikipedia.org/wiki/Wavefront_.obj_file
        """

        self.illumination = illumination

class WaveFrontFileWriter(object):
    def header(self):
        return ""
    def generate_file(self,*args):
        pass

class ObjGraphicFileWriter(WaveFrontFileWriter):
    """Takes multiple Obj objects and writes them to a single OBJ file taking into account
    bookkeeping for vertices and material file association
    """

    def __init__(self,list_of_objects, list_of_materials = []):
        self.list_of_objects = list_of_objects
        self.processed_list_of_objects = []
        self._process_list_of_objects()
        self.list_of_materials = list_of_materials

    def _process_list_of_objects(self):
        for obj in self.list_of_objects:
            if hasattr(obj, "object_graph"): #
                self.processed_list_of_objects.append(obj.object_graph)
            elif obj.__class__.__name__ == "":
                self.process_list_of_objects.append(obj)

    def generate_file(self, file_object, mtl_filename = "material.mtl", mtl_file_object=None):
        "Generates an MTL file"
        f = file_object
        f.write(self.header())

        vertex_index_start = 0
        texture_coordinates_index_start = 0
        normal_vector_index_start = 0

        material_list = self.list_of_materials
        # get materials for processing
        for obj in self.processed_list_of_objects:
            if obj.material_object:
                material_list.append(obj.material_object)

        if len(material_list):
            f.write("mtllib %s\n" % mtl_filename)

        obj_count = 0
        for obj in self.processed_list_of_objects:
            number_of_vertices_added = 0
            matrix_dimension = obj.vertices.shape

            for i in range(matrix_dimension[0]):
                f.write("v")
                for j in range(matrix_dimension[1]):
                    f.write(" " + str(obj.vertices[i,j]))

                number_of_vertices_added += 1
                f.write("\n")

            if obj.has_texture_coordinates:
                tc_matrix_dimension = obj.texture_coordinates.shape
                number_of_texture_coordinated_added = 0
                for i in range(tc_matrix_dimension[0]):
                    f.write("vt")
                    for j in range(tc_matrix_dimension[1]):
                        f.write(" " + str(obj.texture_coordinates[i,j]))
                    number_of_texture_coordinated_added += 1
                    f.write("\n")

            if obj.has_normal_vectors:
                nv_matrix_dimension = obj.normal_vectors.shape
                number_of_normal_vectors_added = 0
                for i in range(nv_matrix_dimension[0]):
                    f.write("vn")
                    for j in range(nv_matrix_dimension[1]):
                        f.write(" " + str(obj.normal_vectors[i,j]))
                    number_of_normal_vectors_added += 1
                    f.write("\n")

            #Group
            f.write("g obj%s\n" % obj_count)
            #assign material
            if obj.material_object:
                f.write("usemtl %s\n" % obj.material_object.name())
            else:
                f.write("usemtl\n")

            #Generate faces
            if obj.has_texture_indices and obj.has_normal_indices:
                format = "ftn" #f 1/2/3 2/2/3
            else:
                if obj.has_texture_indices:
                    format = "ft" #Vertex/texture-coordinate:   f v1/vt1 v2/vt2 v3/vt3 ...
                elif obj.has_normal_indices:
                    format = "fn" # f 1//2 3//2
                else:
                    format = "f" #f 1 2 3

            matrix_dimension = obj.vertices_indices.shape
            for i in range(matrix_dimension[0]): #TODO: add other vertices format
                f.write("f")
                for j in range(matrix_dimension[1]):
                    if format == "f":
                        f.write(" %s" % (obj.vertices_indices[i,j] + vertex_index_start,))
                    elif format  == "ft":
                        f.write(" %s/%s" % (obj.vertices_indices[i,j] + vertex_index_start,obj.texture_indices[i,j] + texture_coordinates_index_start))
                    elif format == "fn":
                        f.write(" %s//%s" % (obj.vertices_indices[i,j] + vertex_index_start, obj.normal_indices[i,j] + normal_vector_index_start))
                    elif format == "ftn":
                        f.write(" %s/%s/%s" % (obj.vertices_indices[i,j] + vertex_index_start, obj.texture_indices[i,j] + texture_coordinates_index_start,obj. normal_indices[i,j] + normal_vector_index_start))
                f.write("\n")
            vertex_index_start += number_of_vertices_added

            if obj.has_texture_coordinates:
                texture_coordinates_index_start += number_of_texture_coordinated_added

            if obj.has_normal_vectors:
                normal_vector_index_start += number_of_normal_vectors_added
            obj_count += 1

        if len(material_list):
            mtlf_obj = MtlFileWriter(material_list)
            if mtl_file_object:
                mtlf_obj.generate_file(mtl_file_object)

    def header(self):
        return "# Wavefront compatible obj file. See: http://en.wikipedia.org/wiki/Wavefront_.obj_file\n"

#
#newmtl Textured
#Ka 1.000 1.000 1.000
#Kd 1.000 1.000 1.000
#Ks 0.000 0.000 0.000
#d 1.0
#illum 2
#map_Ka lenna.tga           # the ambient texture map
#map_Kd lenna.tga           # the diffuse texture map (most of the time, it will
## be the same as the ambient texture map)
#map_Ks lenna.tga           # specular color texture map
#map_Ns lenna_spec.tga      # specular highlight component
#map_d lenna_alpha.tga      # the alpha texture map
#map_bump lenna_bump.tga    # some implementations use 'map_bump' instead of 'bump' below

class MtlFileWriter(WaveFrontFileWriter):
    """Takes multiple MtlFile objects and creates a single MTL File"""

    def __init__(self,list_of_materials):
        self.list_of_mtls = []
        for material in list_of_materials:
            if material.__class__.__name__ == "Material":
                mtl_obj = material.mtl_material_obj
            elif material.__class__.__name__ == "NamedMaterial":
                mtl_obj = None
            else:
                mtl_obj = material
            if mtl_obj:
                self.list_of_mtls.append(mtl_obj)

    def generate_file(self,file_object):
        "Writes material file"
        f = file_object
        f.write("#MTL material file\n")

        for mtl_obj in self.list_of_mtls:
            f.write("newmtl %s\n" % mtl_obj.name)
            if mtl_obj.diffuse_color:
                f.write("Kd %s %s %s\n" % tuple(mtl_obj.diffuse_color))

            if mtl_obj.map_diffuse :
                f.write("map_Kd %s\n" % mtl_obj.map_diffuse)

            if mtl_obj.specular_color:
                f.write("Ks %s %s %s\n" % tuple(mtl_obj.specular_color))

            #TODO: Add more material options